package org.ehe.ai.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.ehe.ai.domain.dto.ChatRequest;
import org.ehe.ai.service.ChatService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.net.URI;
import java.util.concurrent.*;

/**
 * WebSocket处理器
 * @author :zhangnn
 * @className :ChatWebSocketHandler
 * @description: TODO
 * @date 2025-08-23 09:55:16
 */

@Slf4j
@Component
public class ChatWebSocketHandler implements WebSocketHandler {
    @Autowired
    private ChatService chatService;

    private final ObjectMapper objectMapper = new ObjectMapper();

    // 存储活跃的WebSocket连接
    private final ConcurrentHashMap<String, WebSocketSession> activeSessions = new ConcurrentHashMap<>();

    // 线程池
    private final ScheduledExecutorService executorService = Executors.newScheduledThreadPool(10);

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = extractSessionId(session);
        log.info("WebSocket连接建立: {}", sessionId);

        activeSessions.put(sessionId, session);

        // 发送连接成功消息
        sendMessage(session, "connected", "WebSocket连接建立成功，会话ID: " + sessionId, null);

        // 启动心跳
        startHeartbeat(sessionId, session);
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String sessionId = extractSessionId(session);
        log.info("收到WebSocket消息，会话ID: {}, 消息: {}", sessionId, message.getPayload());

        try {
            // 解析消息
            String payload = message.getPayload().toString();
            ChatRequest request = objectMapper.readValue(payload, ChatRequest.class);
            request.setSessionId(sessionId);

            // 异步处理消息
            CompletableFuture.runAsync(() -> processMessage(session, request), executorService);

        } catch (Exception e) {
            log.error("处理WebSocket消息失败", e);
            sendMessage(session, "error", "消息处理失败: " + e.getMessage(), null);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        String sessionId = extractSessionId(session);
        log.error("WebSocket传输错误，会话ID: {}", sessionId, exception);
        activeSessions.remove(sessionId);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String sessionId = extractSessionId(session);
        log.info("WebSocket连接关闭，会话ID: {}, 状态: {}", sessionId, closeStatus);
        activeSessions.remove(sessionId);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 处理聊天消息
     */
    private void processMessage(WebSocketSession session, ChatRequest request) {
        try {
            // 发送开始处理消息
            sendMessage(session, "start", "开始处理您的问题...", null);
            // 调用AI API并流式返回
            String aiResponse = callAiAndStreamResponse(session, request.getMessage());
            // 保存消息
            chatService.saveMessage(request.getSessionId(), request.getMessage(), 1);
            chatService.saveMessage(request.getSessionId(), aiResponse, 2);
            // 发送完成消息
            sendMessage(session, "complete", "回答完成", aiResponse);

        } catch (Exception e) {
            log.error("处理消息失败", e);
            sendMessage(session, "error", "处理失败: " + e.getMessage(), null);
        }
    }

    /**
     * 调用AI API并流式返回响应
     */
    private String callAiAndStreamResponse(WebSocketSession session, String prompt) {
        try {
            String fullResponse = chatService.callAiApi(prompt);
            // 流式发送响应
            streamResponse(session, fullResponse);
            return fullResponse;
        } catch (Exception e) {
            log.error("AI API调用失败", e);
            return "抱歉，AI服务暂时不可用。";
        }
    }

    /**
     * 流式发送响应
     */
    private void streamResponse(WebSocketSession session, String fullResponse) {
        try {
            // 将响应分块发送
            String[] chunks = splitIntoChunks(fullResponse, 15);

            for (String chunk : chunks) {
                sendMessage(session, "chunk", chunk, null);
                Thread.sleep(150 + (int)(Math.random() * 100)); // 模拟打字效果
            }

        } catch (Exception e) {
            log.error("流式发送失败", e);
        }
    }

    /**
     * 分割文本为块
     */
    private String[] splitIntoChunks(String text, int chunkSize) {
        return text.split("(?<=\\G.{" + chunkSize + "})");
    }

    /**
     * 发送WebSocket消息
     */
    private void sendMessage(WebSocketSession session, String type, String message, String fullContent) {
        try {
            if (session.isOpen()) {
                ChatWebSocketMessage wsMessage = new ChatWebSocketMessage();
                wsMessage.setType(type);
                wsMessage.setMessage(message);
                wsMessage.setFullContent(fullContent);
                wsMessage.setTimestamp(System.currentTimeMillis());

                String jsonMessage = objectMapper.writeValueAsString(wsMessage);
                session.sendMessage(new TextMessage(jsonMessage));
            }
        } catch (IOException e) {
            log.error("发送WebSocket消息失败", e);
        }
    }

    /**
     * 启动心跳
     */
    private void startHeartbeat(String sessionId, WebSocketSession session) {
        executorService.scheduleAtFixedRate(() -> {
            if (activeSessions.containsKey(sessionId) && session.isOpen()) {
                sendMessage(session, "heartbeat", "ping", null);
            } else {
                activeSessions.remove(sessionId);
            }
        }, 30, 30, TimeUnit.SECONDS);
    }

    /**
     * 从WebSocket会话中提取sessionId
     */
    private String extractSessionId(WebSocketSession session) {
        URI uri = session.getUri();
        if (uri != null) {
            String path = uri.getPath();
            String[] segments = path.split("/");
            return segments[segments.length - 1];
        }
        return "unknown";
    }

    /**
     * 广播消息到指定会话
     */
    public void broadcastToSession(String sessionId, String type, String message) {
        WebSocketSession session = activeSessions.get(sessionId);
        if (session != null && session.isOpen()) {
            sendMessage(session, type, message, null);
        }
    }

    /**
     * 获取活跃连接数
     */
    public int getActiveConnectionCount() {
        return activeSessions.size();
    }

    /**
     * WebSocket消息对象
     */
    public static class ChatWebSocketMessage {
        private String type;
        private String message;
        private String fullContent;
        private long timestamp;

        // Getters and Setters
        public String getType() { return type; }
        public void setType(String type) { this.type = type; }

        public String getMessage() { return message; }
        public void setMessage(String message) { this.message = message; }

        public String getFullContent() { return fullContent; }
        public void setFullContent(String fullContent) { this.fullContent = fullContent; }

        public long getTimestamp() { return timestamp; }
        public void setTimestamp(long timestamp) { this.timestamp = timestamp; }
    }
}
